/*-------------------------------------------------------------------
mat shell

Ax = b

已知x, 求b

-------------------------------------------------------------------*/

#include <petscmat.h>
#include <time.h>
#include <stdlib.h>

static char help[] = "mat shell\n\n";
extern PetscErrorCode poissonMult(Mat A, Vec x, Vec y);

typedef struct
{
    PetscReal   param;
    PetscInt    gridNum;

    PetscReal   x_start;
    PetscReal   x_end;
    PetscReal   y_start;
    PetscReal   y_end;

    Vec         x;
    Vec         b;

} userData;

int main(int argc, char **args)
{
	PetscErrorCode 	ierr;
    userData        user;
	Vec            	x, b;
	Mat				A;
	PetscInt		gridNum = 4, its;
    PetscInt        localSize = 0;


    // 初始化
    ierr = PetscInitialize(&argc, &args, (char*)0, help); if (ierr) return ierr;
    
    // 命令行参数
    ierr = PetscOptionsGetInt(NULL, NULL, "-gridnum", &gridNum, NULL); CHKERRQ(ierr);

    // 格点数量
    PetscInt gridPointNum = gridNum*gridNum;
    user.param = 1.0;
    user.gridNum = gridNum;
    user.x_start = -1.0;
    user.x_end = 1.0;
    user.y_start = -1.0;
    user.y_end = 1.0;

    // x
    ierr = VecCreate(PETSC_COMM_WORLD, &x); CHKERRQ(ierr);
    ierr = VecSetSizes(x, PETSC_DECIDE, gridPointNum); CHKERRQ(ierr);
    ierr = VecSetFromOptions(x); CHKERRQ(ierr);
    ierr = VecDuplicate(x, &b); CHKERRQ(ierr);
    ierr = VecSet(x, 1.0); CHKERRQ(ierr);

    // A
    ierr = VecGetLocalSize(x, &localSize); CHKERRQ(ierr);
    ierr = MatCreateShell(PETSC_COMM_WORLD, localSize, localSize, gridPointNum, gridPointNum, (void *)&user, &A); CHKERRQ(ierr);
    ierr = MatShellSetOperation(A, MATOP_MULT, (void(*)(void))poissonMult); CHKERRQ(ierr);
    ierr = MatSetFromOptions(A); CHKERRQ(ierr);

    ierr = MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY); CHKERRQ(ierr);
    ierr = MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY); CHKERRQ(ierr);

    // 
    ierr = MatMult(A, x, b); CHKERRQ(ierr);

	// 打印A, x, b
	printf("\nA:\n");
	ierr = MatView(A, PETSC_VIEWER_STDOUT_WORLD); CHKERRQ(ierr);\
	printf("\nx:\n");
	ierr = VecView(x, PETSC_VIEWER_STDOUT_WORLD); CHKERRQ(ierr);
	printf("\nb:\n");
	ierr = VecView(b, PETSC_VIEWER_STDOUT_WORLD); CHKERRQ(ierr);

    // 释放内存
    ierr = VecDestroy(&x); CHKERRQ(ierr);
    ierr = VecDestroy(&b); CHKERRQ(ierr);
    ierr = MatDestroy(&A); CHKERRQ(ierr);

    ierr = PetscFinalize();
    return ierr;
}

// 自定义矩阵乘法
// y = Ax
PetscErrorCode poissonMult(Mat A, Vec x, Vec y)
{
    PetscErrorCode      ierr;
    userData            *puser;
    const PetscScalar   *px;
    PetscScalar         *py;
    PetscInt            n, i;

    PetscFunctionBeginUser;
    MatShellGetContext(A, (void *)puser);

    //ierr = VecCopy(puser->x, x); CHKERRQ(ierr);
    ierr = VecGetArrayRead(x, &px); CHKERRQ(ierr);
    ierr = VecGetArray(y, &py); CHKERRQ(ierr);

    ierr = VecGetSize(x, &n); CHKERRQ(ierr);

    for (i = 0; i < n; i++)
    {
        py[i] = px[i] + 2;
    }

    ierr = VecRestoreArrayRead(x, &px); CHKERRQ(ierr);
    ierr = VecRestoreArray(y, &py); CHKERRQ(ierr);

    PetscFunctionReturn(0);
}